package com.hunttown.common.sharding.client.mybatis;

import com.hunttown.common.sharding.client.annotation.ISharding;
import com.hunttown.common.sharding.client.annotation.Sharding;
import com.hunttown.common.sharding.client.datasource.ShardingDataSources;
import com.hunttown.common.sharding.client.util.BeanPropertyAccessUtil;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import javax.sql.DataSource;

import org.apache.ibatis.exceptions.PersistenceException;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.reflection.ExceptionUtil;
import org.apache.ibatis.session.*;
import org.apache.ibatis.session.defaults.DefaultSqlSession;
import org.apache.ibatis.transaction.Transaction;
import org.apache.ibatis.transaction.TransactionFactory;
import org.mybatis.spring.MyBatisExceptionTranslator;
import org.mybatis.spring.SqlSessionHolder;
import org.mybatis.spring.SqlSessionTemplate;
import org.mybatis.spring.SqlSessionUtils;
import org.mybatis.spring.transaction.SpringManagedTransactionFactory;
import org.springframework.dao.TransientDataAccessResourceException;
import org.springframework.dao.support.PersistenceExceptionTranslator;
import org.springframework.jdbc.CannotGetJdbcConnectionException;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.jdbc.datasource.TransactionAwareDataSourceProxy;
import org.springframework.transaction.support.TransactionSynchronizationAdapter;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.Assert;

public class ShardSqlSessionTemplate extends SqlSessionTemplate {
    private final SqlSessionFactory sqlSessionFactory;
    private final ExecutorType executorType;
    private final SqlSession sqlSessionProxy;
    private final PersistenceExceptionTranslator exceptionTranslator;

    public ShardSqlSessionTemplate(SqlSessionFactory sqlSessionFactory) {
        this(sqlSessionFactory, sqlSessionFactory.getConfiguration().getDefaultExecutorType());
    }

    public ShardSqlSessionTemplate(SqlSessionFactory sqlSessionFactory, ExecutorType executorType) {
        this(sqlSessionFactory, executorType, new MyBatisExceptionTranslator(sqlSessionFactory.getConfiguration().getEnvironment().getDataSource(), true));
    }

    public ShardSqlSessionTemplate(SqlSessionFactory sqlSessionFactory, ExecutorType executorType, PersistenceExceptionTranslator exceptionTranslator) {
        super(sqlSessionFactory, executorType, exceptionTranslator);
        Assert.notNull(sqlSessionFactory, "Property 'sqlSessionFactory' is required");
        Assert.notNull(executorType, "Property 'executorType' is required");
        this.sqlSessionFactory = sqlSessionFactory;
        this.executorType = executorType;
        this.exceptionTranslator = exceptionTranslator;
        this.sqlSessionProxy = (SqlSession) Proxy.newProxyInstance(SqlSessionFactory.class.getClassLoader(), new Class[]{SqlSession.class}, new ShardSqlSessionTemplate.SqlSessionInterceptor());
    }

    public SqlSessionFactory getSqlSessionFactory() {
        return this.sqlSessionFactory;
    }

    public ExecutorType getExecutorType() {
        return this.executorType;
    }

    public PersistenceExceptionTranslator getPersistenceExceptionTranslator() {
        return this.exceptionTranslator;
    }

    public Object selectOne(String statement) {
        return this.sqlSessionProxy.selectOne(statement);
    }

    public Object selectOne(String statement, Object parameter) {
        return this.sqlSessionProxy.selectOne(statement, parameter);
    }

    public Map<?, ?> selectMap(String statement, String mapKey) {
        return this.sqlSessionProxy.selectMap(statement, mapKey);
    }

    public Map<?, ?> selectMap(String statement, Object parameter, String mapKey) {
        return this.sqlSessionProxy.selectMap(statement, parameter, mapKey);
    }

    public Map<?, ?> selectMap(String statement, Object parameter, String mapKey, RowBounds rowBounds) {
        return this.sqlSessionProxy.selectMap(statement, parameter, mapKey, rowBounds);
    }

    public List<?> selectList(String statement) {
        return this.sqlSessionProxy.selectList(statement);
    }

    public List<?> selectList(String statement, Object parameter) {
        return this.sqlSessionProxy.selectList(statement, parameter);
    }

    public List<?> selectList(String statement, Object parameter, RowBounds rowBounds) {
        return this.sqlSessionProxy.selectList(statement, parameter, rowBounds);
    }

    public void select(String statement, ResultHandler handler) {
        this.sqlSessionProxy.select(statement, handler);
    }

    public void select(String statement, Object parameter, ResultHandler handler) {
        this.sqlSessionProxy.select(statement, parameter, handler);
    }

    public void select(String statement, Object parameter, RowBounds rowBounds, ResultHandler handler) {
        this.sqlSessionProxy.select(statement, parameter, rowBounds, handler);
    }

    public int insert(String statement) {
        return this.sqlSessionProxy.insert(statement);
    }

    public int insert(String statement, Object parameter) {
        return this.sqlSessionProxy.insert(statement, parameter);
    }

    public int update(String statement) {
        return this.sqlSessionProxy.update(statement);
    }

    public int update(String statement, Object parameter) {
        return this.sqlSessionProxy.update(statement, parameter);
    }

    public int delete(String statement) {
        return this.sqlSessionProxy.delete(statement);
    }

    public int delete(String statement, Object parameter) {
        return this.sqlSessionProxy.delete(statement, parameter);
    }

    public <T> T getMapper(Class<T> type) {
        return this.getConfiguration().getMapper(type, this);
    }

    public void commit() {
        throw new UnsupportedOperationException("Manual commit is not allowed over a Spring managed SqlSession");
    }

    public void commit(boolean force) {
        throw new UnsupportedOperationException("Manual commit is not allowed over a Spring managed SqlSession");
    }

    public void rollback() {
        throw new UnsupportedOperationException("Manual rollback is not allowed over a Spring managed SqlSession");
    }

    public void rollback(boolean force) {
        throw new UnsupportedOperationException("Manual rollback is not allowed over a Spring managed SqlSession");
    }

    public void close() {
        throw new UnsupportedOperationException("Manual close is not allowed over a Spring managed SqlSession");
    }

    public void clearCache() {
        this.sqlSessionProxy.clearCache();
    }

    public Configuration getConfiguration() {
        return this.sqlSessionFactory.getConfiguration();
    }

    public Connection getConnection() {
        return this.sqlSessionProxy.getConnection();
    }

    private class SqlSessionInterceptor implements InvocationHandler {
        private SqlSessionInterceptor() {
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            SqlSession sqlSession = this.getSqlSession(this.getRoutedSingleDataSource(args), ShardSqlSessionTemplate.this.executorType, ShardSqlSessionTemplate.this.exceptionTranslator);

            Object unwrapped;
            try {
                Object result = method.invoke(sqlSession, args);
                if (!SqlSessionUtils.isSqlSessionTransactional(sqlSession, ShardSqlSessionTemplate.this.sqlSessionFactory)) {
                    sqlSession.commit();
                }

                unwrapped = result;
            } catch (Throwable var10) {
                unwrapped = ExceptionUtil.unwrapThrowable(var10);
                if (ShardSqlSessionTemplate.this.exceptionTranslator != null && unwrapped instanceof PersistenceException) {
                    unwrapped = ShardSqlSessionTemplate.this.exceptionTranslator.translateExceptionIfPossible((PersistenceException) unwrapped);
                }

                throw (Throwable) unwrapped;
            } finally {
                SqlSessionUtils.closeSqlSession(sqlSession, ShardSqlSessionTemplate.this.sqlSessionFactory);
            }

            return unwrapped;
        }

        public SqlSession getSqlSession(DataSource dataSource, ExecutorType executorType, PersistenceExceptionTranslator exceptionTranslator) {
            Assert.notNull(ShardSqlSessionTemplate.this.sqlSessionFactory, "No SqlSessionFactory specified");
            Assert.notNull(executorType, "No ExecutorType specified");
            SqlSessionHolder holder = (SqlSessionHolder) TransactionSynchronizationManager.getResource(ShardSqlSessionTemplate.this.sqlSessionFactory);
            if (holder != null && holder.isSynchronizedWithTransaction()) {
                if (holder.getExecutorType() != executorType) {
                    throw new TransientDataAccessResourceException("Cannot change the ExecutorType when there is an existing transaction");
                } else {
                    holder.requested();
                    return holder.getSqlSession();
                }
            } else {
                boolean transactionAware = dataSource instanceof TransactionAwareDataSourceProxy;

                Connection conn;
                try {
                    conn = transactionAware ? dataSource.getConnection() : DataSourceUtils.getConnection(dataSource);
                } catch (SQLException var11) {
                    throw new CannotGetJdbcConnectionException("Could not get JDBC Connection for SqlSession", var11);
                }

                TransactionFactory transactionFactory = new SpringManagedTransactionFactory();
                Transaction tx = transactionFactory.newTransaction(dataSource, TransactionIsolationLevel.REPEATABLE_READ, false);
                Executor executor = ShardSqlSessionTemplate.this.getConfiguration().newExecutor(tx, executorType);
                SqlSession session = new DefaultSqlSession(ShardSqlSessionTemplate.this.getConfiguration(), executor, false);
                if (TransactionSynchronizationManager.isSynchronizationActive()) {
                    if (!(ShardSqlSessionTemplate.this.getConfiguration().getEnvironment().getTransactionFactory() instanceof SpringManagedTransactionFactory) && DataSourceUtils.isConnectionTransactional(conn, dataSource)) {
                        throw new TransientDataAccessResourceException("SqlSessionFactory must be using a SpringManagedTransactionFactory in order to use Spring transaction synchronization");
                    }

                    holder = new SqlSessionHolder(session, executorType, exceptionTranslator);
                    TransactionSynchronizationManager.bindResource(ShardSqlSessionTemplate.this.sqlSessionFactory, holder);
                    TransactionSynchronizationManager.registerSynchronization(new ShardSqlSessionTemplate.SqlSessionInterceptor.SqlSessionSynchronization(holder, ShardSqlSessionTemplate.this.sqlSessionFactory));
                    holder.setSynchronizedWithTransaction(true);
                    holder.requested();
                }

                return session;
            }
        }

        protected DataSource getRoutedSingleDataSource(Object[] parameters) {
            String dataSourceName = null;
            Object[] arr$ = parameters;
            int len$ = parameters.length;

            for (int i$ = 0; i$ < len$; ++i$) {
                Object parameterObj = arr$[i$];

                try {
                    if (parameterObj instanceof List) {
                        List<Object> shardList = (List) parameterObj;
                        if (shardList != null && shardList.size() > 0) {
                            Object shardObj = shardList.get(0);
                            if (shardObj instanceof Sharding || shardObj instanceof ISharding) {
                                dataSourceName = (String) BeanPropertyAccessUtil.getPropertyValue("dataSourceName", shardObj);
                            }
                        }
                    } else if (parameterObj instanceof Map) {
                        Map<String, Object> shardMap = (Map) parameterObj;
                        if (shardMap.containsKey("dataSourceName")) {
                            dataSourceName = (String) shardMap.get("dataSourceName");
                        }
                    } else if (parameterObj instanceof Sharding || parameterObj instanceof ISharding) {
                        dataSourceName = (String) BeanPropertyAccessUtil.getPropertyValue("dataSourceName", parameterObj);
                    }
                } catch (Throwable var9) {
                    var9.printStackTrace();
                }

                if (dataSourceName != null) {
                    break;
                }
            }

            Assert.notNull(dataSourceName, "get DataSourceName error!");
            return dataSourceName != null ? this.getDataSourceByName(dataSourceName) : null;
        }

        public DataSource getDataSourceByName(String dataSourceName) {
            return ShardingDataSources.getInstance().getDataSource(dataSourceName);
        }

        private final class SqlSessionSynchronization extends TransactionSynchronizationAdapter {
            private final SqlSessionHolder holder;
            private final SqlSessionFactory sessionFactory;

            public SqlSessionSynchronization(SqlSessionHolder holder, SqlSessionFactory sessionFactory) {
                Assert.notNull(holder, "Parameter 'holder' must be not null");
                Assert.notNull(sessionFactory, "Parameter 'sessionFactory' must be not null");
                this.holder = holder;
                this.sessionFactory = sessionFactory;
            }

            public int getOrder() {
                return 999;
            }

            public void suspend() {
                TransactionSynchronizationManager.unbindResource(this.sessionFactory);
            }

            public void resume() {
                TransactionSynchronizationManager.bindResource(this.sessionFactory, this.holder);
            }

            public void beforeCommit(boolean readOnly) {
                if (TransactionSynchronizationManager.isActualTransactionActive()) {
                    try {
                        this.holder.getSqlSession().commit();
                    } catch (PersistenceException var3) {
                        if (this.holder.getPersistenceExceptionTranslator() != null) {
                            throw this.holder.getPersistenceExceptionTranslator().translateExceptionIfPossible(var3);
                        }

                        throw var3;
                    }
                }

            }

            public void afterCompletion(int status) {
                if (!this.holder.isOpen()) {
                    TransactionSynchronizationManager.unbindResource(this.sessionFactory);

                    try {
                        this.holder.getSqlSession().close();
                    } finally {
                        this.holder.reset();
                    }
                }

            }
        }
    }
}